import pandas as pd
import numpy as np
from math import ceil
from datetime import datetime

# 加载与预处理
def load():
    lines = pd.read_excel('./rawData/附件1.xlsx')    # 线路信息
    owners = pd.read_excel('./rawData/附件5.xlsx')   # 自有车数
    routes = pd.read_excel('./rawData/附件4.xlsx')   # 串点站点
    preds  = pd.read_excel('./processedData/结果表1_预测结果.xlsx')  # 预测结果

    # 时间与日期
    lines['time'] = pd.to_datetime(lines['发运节点'], format='%H:%M:%S')
    lines['date'] = '2024/12/16'
    preds['date'] = pd.to_datetime(preds['日期']).dt.strftime('%Y/%m/%d')

    # 编码清洗
    lines['code'] = lines['线路编码'].str.strip()
    preds['code']  = preds['线路编码'].str.strip()

    # 预测量映射字典
    vol_map = preds.set_index(['code','date'])['货量'].to_dict()
    lines['vol'] = lines.apply(
        lambda r: vol_map.get((r['code'], r['date']), 0),
        axis=1
    )

    return lines, owners, routes

# 构建双向串点集合
def pairs(routes):
    s1 = {(r['站点编号1'], r['站点编号2']) for _, r in routes.iterrows()}
    s2 = {(r['站点编号2'], r['站点编号1']) for _, r in routes.iterrows()}
    return s1 | s2

# 容器决策与需求计算
def decide(lines, C0=1000, C1=800):
    lines['hr'] = lines['time'].dt.hour
    # 条件：货量 ≤ C1 或 编码以 Z 开头 或 6 点前
    lines['cont'] = np.where(
        (lines['vol'] <= C1) |
        (lines['车队编码'].str.startswith('Z')) |
        (lines['hr'] <= 6),
        1, 0
    )
    # 单车容量 & 需求车数
    lines['cap']  = np.where(lines['cont']==1, C1, C0)
    lines['need'] = np.ceil(lines['vol'] / lines['cap']).astype(int)
    return lines

# 判断两条线能否串联
def can_chain(a, b, prs):
    dt = abs((a['time'] - b['time']).total_seconds())
    return (
        a['起始场地']==b['起始场地'] and
        a['车队编码']==b['车队编码'] and
        dt <= 1800 and
        (a['目的场地']==b['目的场地'] or (a['目的场地'], b['目的场地']) in prs) and
        a['cont']==b['cont']
    )

# 贪心串点
def chain(lines, prs):
    used = set()
    out  = []
    df   = lines.sort_values('time')
    for _, x in df.iterrows():
        if x['code'] in used: continue
        grp = [x]
        used.add(x['code'])
        vol_sum = x['vol']
        cap     = x['cap']
        for _, y in df.iterrows():
            if y['code'] in used: continue
            if can_chain(x, y, prs) and vol_sum + y['vol'] <= cap:
                grp.append(y)
                used.add(y['code'])
                vol_sum += y['vol']
        out.append(grp)
    return out

# 分配车辆
def assign(chains, owners):
    pool = {t:c for t,c in zip(owners['车队编码'], owners['自有车数量'])}
    res  = []
    vid  = 0
    for grp in sorted(chains, key=lambda g: -len(g)):
        team = grp[0]['车队编码']
        typ  = '自有' if pool.get(team,0)>0 else '外部'
        if typ=='自有': pool[team] -= 1
        vid += 1
        for r in grp:
            res.append({
                'code': r['code'],
                'date': r['date'],
                'time': r['time'].strftime('%H:%M:%S'),
                'cont': '是' if r['cont'] else '否',
                'veh': f'{typ}-V{vid}'
            })
    return pd.DataFrame(res)

# 主流程
def main():
    lines, owners, routes = load()
    prs    = pairs(routes)
    lines  = decide(lines)
    chs    = chain(lines, prs)
    plan   = assign(chs, owners)
    plan.to_excel('./processedData/结果表4.xlsx', index=False)
    print('调度结果已写入：结果表4.xlsx')
    tgt = ['场地3 - 站点83 - 0600','场地3 - 站点83 - 1400']
    print('\n指定线路：\n', plan[plan['code'].isin(tgt)])

if __name__ == '__main__':
    main()
